import os
import PIL.Image as Image
import torch
from torch.utils.data import Dataset
import re
from denoising_config import *

#定义函数，对图片排序
def sort_alphanum(img_names):
    #数字转Int
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    aplhanum_key = lambda img_names: [convert(c) for c in re.split('([0-9]+)', img_names)]
    return sorted(img_names, key=aplhanum_key)

class ImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        super().__init__()
        self.image_dir = image_dir
        self.transform = transform
        self.all_images = sort_alphanum(os.listdir(image_dir))

    def __len__(self):
        return len(self.all_images)

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.all_images[idx])
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        noise_factor = 0.5
        noisy_images = image + noise_factor*torch.randn(*image.shape)
        noisy_images = torch.clamp(noisy_images, 0., 1.)
        #返回（噪声图片，原始图片）
        return noisy_images,image
    


if __name__ == "__main__":
    images = os.listdir(IMG_PATH)
    sorted_images = sort_alphanum(images)
    # print(sorted_images)
